package com.tuniu.agents.memory;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * Chat Memory Manager
 */
public class ChatMemoryManager implements MemoryManager{

    private static final Logger logger = LoggerFactory.getLogger(ChatMemoryManager.class);

    private final ChatMemory chatMemory;

    public ChatMemoryManager(ChatMemory chatMemory) {
        this.chatMemory = chatMemory;
    }

    /**
     * Get all contexts under the session ID
     * @param conversationId
     * @return List of Messages
     */
    @Override
    public List<Message> getChatMessage(String conversationId) {
        return chatMemory.get(conversationId,Integer.MAX_VALUE);
    }

    /**
     * Replace all context information
     * @param conversationId
     * @param messages
     * @return boolean
     */
    @Override
    public boolean setMessageAll(String conversationId, List<Message> messages) {
        try {
            if(!isAgentIdExist(messages)){
                return false;
            }
            chatMemory.clear(conversationId);
            chatMemory.add(conversationId, messages);
        }catch (Exception e){
            logger.error("ChatMemoryManager.setMessageAll e:",e);
            return false;
        }
        return true;
    }

    /**
     * Get the context of the specified agentId under the specified session ID
     * @param conversationId
     * @param agentId
     * @return List of Messages
     */
    @Override
    public List<Message> getAgentMessage(String conversationId, String agentId) {
        List<Message> agentMessages = new ArrayList<>();
        List<Message> chatMessages = chatMemory.get(conversationId,Integer.MAX_VALUE);
        // Get messages specific to agentId
        chatMessages.forEach(message -> {
            if(message.getMetadata().get("agentId").toString().equals(agentId)){
                agentMessages.add(message);
            }
        });
        return agentMessages;
    }

    @Override
    public boolean setMessageRight(String conversationId, List<Message> messages) {
        try {
            if(!isAgentIdExist(messages)){
                return false;
            }
            chatMemory.add(conversationId, messages);
        }catch (Exception e){
            logger.error("ChatMemoryManager.setMessageRight e:",e);
            return false;
        }
        return true;
    }

    /**
     * Insert chatMessage at the specified round position of the session ID
     * @param conversationId
     * @param message
     * @param round
     * @return boolean
     */
    //@Override
    public boolean setRoundMessage(String conversationId, List<Message> message, int round) {
        try {
            if(!isAgentIdExist(message)){
                return false;
            }
            List<Message> chatMessages = chatMemory.get(conversationId, Integer.MAX_VALUE);
            // Replace records at the specified round
            RoundChatMemory roundChatMemory = getRoundChatMemory(chatMessages,round);
            if( roundChatMemory.startIndex!=0 && roundChatMemory.endIndex!=0) {
                chatMessages.subList(roundChatMemory.startIndex, roundChatMemory.endIndex).clear();
            }
            message.forEach(msg -> {msg.getMetadata().put("round",round);});
            chatMessages.addAll(message);
            chatMemory.clear(conversationId);
            chatMemory.add(conversationId, chatMessages);
        }catch (Exception e){
            logger.error("ChatMemoryManager.setMessage e:",e);
            return false;
        }
        return true;
    }

    /**
     * Remove chatMessage of the specified round under the session ID
     * @param conversationId
     * @param round
     * @return boolean
     */
    //@Override
    public boolean removeRound(String conversationId, int round) {
        try {
            List<Message> chatMessages = chatMemory.get(conversationId, Integer.MAX_VALUE);
            RoundChatMemory roundChatMemory = getRoundChatMemory(chatMessages, round);
            chatMessages.subList(roundChatMemory.startIndex, roundChatMemory.endIndex).clear();
            chatMemory.clear(conversationId);
            chatMemory.add(conversationId, chatMessages);
        }catch (Exception e){
            logger.error("ChatMemoryManager.removeRound e:{}",e);
            return false;
        }
        return true;
    }

    @Override
    public boolean remove(String conversationId) {
        try {
            chatMemory.clear(conversationId);
        }catch (Exception e){
            logger.error("ChatMemoryManager.remove e:",e);
            return false;
        }
        return true;
    }

    /**
     * Remove context records of the specified agentId under the specified session ID, keeping only the latest one
     * @param conversationId
     * @param agentId
     * @return boolean
     */
    @Override
    public boolean removeExceptLatest(String conversationId, String agentId) {
        try {
            // Get all messages of the specified session
            List<Message> chatMessages = getChatMessage(conversationId);
            // Filter messages belonging to the specified agentId
            List<Message> agentMessages = getAgentMessage(conversationId,agentId);
            if (agentMessages.isEmpty()) {
                return true;
            }
            Message latestMessage = agentMessages.get(agentMessages.size() - 1);
            chatMessages.removeIf(message -> {
                Object metadataAgentId = message.getMetadata().get("agentId");
                return agentId.equals(metadataAgentId) && !message.equals(latestMessage);
            });
            // Store the modified chatMessages back to Redis
            chatMemory.clear(conversationId);
            chatMemory.add(conversationId, chatMessages);
        }catch (Exception e){
            logger.error("ChatMemoryManager.removeExceptLatest e:",e);
            return false;
        }
        return true;
    }

    /**
     * Remove context records of the specified agentId under the specified session ID
     * @param conversationId
     * @param agentId
     * @return boolean
     */
    @Override
    public boolean removeAgentId(String conversationId, String agentId) {
        try {
            // Get all messages of the specified session
            List<Message> chatMessages = getChatMessage(conversationId);
            chatMessages.removeIf(message -> {
                Object metadataAgentId = message.getMetadata().get("agentId");
                return agentId.equals(metadataAgentId);
            });
            // Store the modified chatMessages back to Redis
            chatMemory.clear(conversationId);
            if(!CollectionUtils.isEmpty(chatMessages)) {
                chatMemory.add(conversationId, chatMessages);
            }
        }catch (Exception e){
            logger.error("ChatMemoryManager.removeAgentId e:",e);
            return false;
        }
        return true;
    }

    private record RoundChatMemory(int startIndex,int endIndex){

    }

    /**
     * Ensure that the metadata and agentId parameters in the new message must exist
     * @param chatMessages
     * @return boolean
     */
    private boolean isAgentIdExist(List<Message> chatMessages){
        for(Message msg : chatMessages){
            if(msg.getMetadata()==null){
                return false;
            }
            if(msg.getMetadata().get("agentId")==null || msg.getMetadata().get("agentId").toString().equals("")){
                return false;
            }
        }
        return true;
    }

    private RoundChatMemory getRoundChatMemory(List<Message> chatMessages,int round){
        int startIndex = 0;
        int endIndex = 0;
        int count = 1;
        for (int i = 0; i < chatMessages.size(); i++) {
            Message msg = chatMessages.get(i);
            if (round == (int) msg.getMetadata().get("round")) {
                if (count == 1) {
                    startIndex = i;
                    count++;
                } else {
                    endIndex = i;
                }
            }
        }
        return new RoundChatMemory(startIndex,endIndex);
    }

}
